from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from matplotlib.ticker import FuncFormatter

def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    # 定义一个函数，用于生成随机权重
    def init_weights_r(key, shape):
        return jax.random.normal(key, shape)
    # 生成一个随机的 PRNGKey
    key = jax.random.PRNGKey(np.random.randint(0, 1000))
    # key = jax.random.PRNGKey(4512)
    # 使用 tree_map 遍历 params 对象，并使用 init_weights 函数生成随机权重
    random_params = jax.tree_map(lambda x: init_weights_r(key, x.shape), params)
    # params = random_params

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])
    mat1 = np.array(tree_leaves[1])
    print("mat1.shape: ", mat1.shape)
    print("bias1: ", bias1)
    bias2 = np.array(tree_leaves[2])
    mat2 = np.array(tree_leaves[3])

    mat_obs = np.array(tree_leaves[1])[rpl_config.nn_size:rpl_config.nn_size+9,:]
    mat_intr = np.array(tree_leaves[1])[0:rpl_config.nn_size,:]
    print("mat_obs.shape: ", mat_obs.shape)

    # 将 mat1 显示为图像，颜色采用默认的 jet 颜色
    plt.imshow(mat_intr, cmap='jet')
    # 显示颜色条
    plt.colorbar()
    # 显示图像
    plt.show()

    # 保存 mat_intr
    np.save("./logs/mat_intr.npy", mat_intr)

if __name__ == "__main__":
    
    ivf()